import torch
import torch_npu
import platform
# 设置npu设备
torch_device = "npu:4" # 0~7，根据自己的设备写
torch.npu.set_device(torch.device(torch_device))
torch.npu.set_compile_mode(jit_compile=False)
option = {}
option["NPU_FUZZY_COMPILE_BLACKLIST"] = "Tril"
torch.npu.set_option(option)

from transformers import AutoModelForCausalLM, AutoTokenizer

# 指定预训练模型的路径，使用时改为正确的模型路径
DEFAULT_CKPT_PATH = '../mistral-7b-instract-v0.3'

# 加载预训练的模型，然后将模型移至NPU设备，并设置为评估模式
model = AutoModelForCausalLM.from_pretrained(
    DEFAULT_CKPT_PATH,
    torch_dtype=torch.float16   
).npu().eval()

# 加载与模型配套的分词器
tokenizer = AutoTokenizer.from_pretrained(DEFAULT_CKPT_PATH)

# 主循环，用于与用户交互
while True:
    prompt = input("user:")
    if prompt == "exit":
        break
    # 准备输入数据，这里构建了一个包含系统提示和用户输入的对话模板
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": prompt}
    ]

    # 使用分词器的apply_chat_template方法将对话模板转换为模型可以接受的格式
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    # 将输入数据转换为PyTorch张量，并移至NPU设备
    model_inputs = tokenizer([text], return_tensors="pt").to(torch_device)
    # 使用模型生成文本
    generated_ids = model.generate(
        model_inputs.input_ids,
        max_new_tokens=512
    )
    # 只保留输入序列之后生成的ID
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]
    # 使用分词器的batch_decode方法将生成的ID解码为文本
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    print("mistral-7b-instract-v0.3:",response)

